import json
from typing import Any, Callable, Iterable, Mapping

from zkl_pyutils_fsspec import FsLike, resolve_fs

from zkl_aiutils_datasets.utils.qual import get_and_check_module_and_qualname
from .loading import dataset_metadata_file_name
from .metadata import DatasetLoaderConfig, DatasetMetadata


def write_dataset_metadata(fs: FsLike, dataset_metadata: DatasetMetadata | Any):
    fs = resolve_fs(fs)
    fs.mkdirs("", exist_ok=True)
    with fs.open(dataset_metadata_file_name, "wt") as fp:
        json.dump(dataset_metadata, fp, indent=2)


def make_dataset_metadata(
    loader: Callable | type | None = None,
    args: Iterable[Any] | None = None,
    kwargs: Mapping[str, Any] | None = None,
) -> DatasetMetadata:
    return DatasetMetadata(loader=make_dataset_loader_config(loader, args, kwargs))


def make_dataset_loader_config(
    loader: Callable | type | None = None,
    args: Iterable[Any] | None = None,
    kwargs: Mapping[str, Any] | None = None,
) -> DatasetLoaderConfig | None:
    if loader is None:
        return None
    module, qualname = get_and_check_module_and_qualname(loader)
    args = tuple(args) if args is not None else None
    kwargs = dict(kwargs) if kwargs is not None else None
    return DatasetLoaderConfig(module=module, qualname=qualname, args=args, kwargs=kwargs)
